/*
 * Copyright (c) 2016, Linaro Limited
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 * this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 * this list of conditions and the following disclaimer in the documentation
 * and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#include <err.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>

/* OP-TEE TEE client API (built by optee_client) */
#include <tee_client_api.h>

/* For the UUID (found in the TA's h-file(s)) */
#include <api_test_ta.h>

#include <openssl/rsa.h>
#include <openssl/pem.h>
#include <openssl/bn.h>
#include <openssl/evp.h>
#include <openssl/engine.h>
#include <assert.h>

void print_hex(uint8_t *data, size_t len)
{
	for (size_t i = 0; i < len; i++)
	{
		printf("%02x", data[i]);
	}
	printf("\n");
}
void AddClient(TEEC_Session *sess, char *path)
{
	printf("Add %s\n", path);
	char *resolved = realpath(path, NULL);
	if (resolved == NULL)
	{
		printf("realpath falied\n");
		return;
	}
	printf("Absolute path: %s\n", resolved);
	TEEC_Result res;
	TEEC_Operation op;
	uint32_t err_origin;

	memset(&op, 0, sizeof(op));

	op.paramTypes = TEEC_PARAM_TYPES(TEEC_MEMREF_TEMP_INPUT, TEEC_NONE,
									 TEEC_NONE, TEEC_NONE);
	op.params[0].tmpref.buffer = resolved;
	op.params[0].tmpref.size = strlen(resolved);

	res = TEEC_InvokeCommand(sess, TA_API_TEST_CMD_ADD, &op, &err_origin);
	free(resolved);
	if (res != TEEC_SUCCESS)
		errx(1, "TEEC_InvokeCommand failed with code 0x%x origin 0x%x", res, err_origin);
}

void DeleteClient(TEEC_Session *sess, char *path)
{

	printf("Delete %s\n", path);
	char *resolved = realpath(path, NULL);
	if (resolved == NULL)
	{
		printf("realpath falied\n");
		return;
	}
	printf("Absolute path: %s\n", resolved);
	TEEC_Result res;
	TEEC_Operation op;
	uint32_t err_origin;

	memset(&op, 0, sizeof(op));

	op.paramTypes = TEEC_PARAM_TYPES(TEEC_MEMREF_TEMP_INPUT, TEEC_NONE,
									 TEEC_NONE, TEEC_NONE);
	op.params[0].tmpref.buffer = resolved;
	op.params[0].tmpref.size = strlen(resolved);

	res = TEEC_InvokeCommand(sess, TA_API_TEST_CMD_DELETE, &op, &err_origin);
	free(resolved);

	if (res != TEEC_SUCCESS)
		errx(1, "TEEC_InvokeCommand failed with code 0x%x origin 0x%x", res, err_origin);
}
#define RSA_SIZE 1024
#define BUFF_SIZE (RSA_SIZE / 8)
uint8_t rsa_key[8][BUFF_SIZE] = {0};
size_t rsa_key_size[8] = {0};
void GetRSA(TEEC_Session *sess)
{
	TEEC_Result res;
	TEEC_Operation op[2] = {0};
	uint32_t err_origin;

	op[0].paramTypes = TEEC_PARAM_TYPES(TEEC_MEMREF_TEMP_OUTPUT, TEEC_MEMREF_TEMP_OUTPUT,
										TEEC_MEMREF_TEMP_OUTPUT, TEEC_MEMREF_TEMP_OUTPUT);

	op[0].params[0].tmpref.buffer = rsa_key[0];
	op[0].params[0].tmpref.size = BUFF_SIZE;
	op[0].params[1].tmpref.buffer = rsa_key[1];
	op[0].params[1].tmpref.size = BUFF_SIZE;
	op[0].params[2].tmpref.buffer = rsa_key[2];
	op[0].params[2].tmpref.size = BUFF_SIZE;
	op[0].params[3].tmpref.buffer = rsa_key[3];
	op[0].params[3].tmpref.size = BUFF_SIZE;

	res = TEEC_InvokeCommand(sess, TA_API_TEST_CMD_RSA_PRI_0, &op[0], &err_origin);
	if (res != TEEC_SUCCESS)
		errx(1, "TEEC_InvokeCommand failed with code 0x%x origin 0x%x", res, err_origin);

	for (int i = 0; i < 4; i++)
	{
		printf("RSA[%d] size : %ld data : %d\n", i, op[0].params[i].tmpref.size, rsa_key[i][0]);
	}

	op[1].paramTypes = TEEC_PARAM_TYPES(TEEC_MEMREF_TEMP_OUTPUT, TEEC_MEMREF_TEMP_OUTPUT,
										TEEC_MEMREF_TEMP_OUTPUT, TEEC_MEMREF_TEMP_OUTPUT);

	op[1].params[0].tmpref.buffer = rsa_key[4];
	op[1].params[0].tmpref.size = BUFF_SIZE;
	op[1].params[1].tmpref.buffer = rsa_key[5];
	op[1].params[1].tmpref.size = BUFF_SIZE;
	op[1].params[2].tmpref.buffer = rsa_key[6];
	op[1].params[2].tmpref.size = BUFF_SIZE;
	op[1].params[3].tmpref.buffer = rsa_key[7];
	op[1].params[3].tmpref.size = BUFF_SIZE / 2;
	res = TEEC_InvokeCommand(sess, TA_API_TEST_CMD_RSA_PRI_1, &op[1], &err_origin);
	if (res != TEEC_SUCCESS)
		errx(1, "TEEC_InvokeCommand failed with code 0x%x origin 0x%x", res, err_origin);
	for (int i = 0; i < 4; i++)
	{
		printf("RSA[%d] size : %ld data : %d\n", i + 4, op[1].params[i].tmpref.size, rsa_key[i + 4][0]);
	}

	for (size_t i = 0; i < 4; i++)
	{
		rsa_key_size[i] = op[0].params[i].tmpref.size;
		rsa_key_size[i + 4] = op[1].params[i].tmpref.size;
	}
	for (size_t i = 0; i < 8; i++)
	{
		print_hex(rsa_key[i], rsa_key_size[i]);
	}
}

enum
{
	RSA_KEY_PART_N = 0,
	RSA_KEY_PART_E,
	RSA_KEY_PART_D,
	RSA_KEY_PART_P,
	RSA_KEY_PART_Q,
	RSA_KEY_PART_DMP1,
	RSA_KEY_PART_DMQ1,
	RSA_KEY_PART_IQMP,
	RSA_KEY_PART_NUM,
};
void print_rsa_key(RSA *rsa)
{
	if (rsa == NULL)
	{
		printf("RSA structure is NULL\n");
		return;
	}

	// 打印公钥
	printf("Public Key (Modulus and Exponent):\n");
	BN_print_fp(stdout, RSA_get0_n(rsa)); // 打印模数 (n)
	printf("\n");
	BN_print_fp(stdout, RSA_get0_e(rsa)); // 打印公钥指数 (e)
	printf("\n");

	// 打印私钥
	printf("Private Key (Exponent, etc.):\n");
	BN_print_fp(stdout, RSA_get0_d(rsa)); // 打印私钥指数 (d)
	printf("\n");
	BN_print_fp(stdout, RSA_get0_p(rsa)); // 打印质因数 p
	printf("\n");
	BN_print_fp(stdout, RSA_get0_q(rsa)); // 打印质因数 q
	printf("\n");
	BN_print_fp(stdout, RSA_get0_dmp1(rsa)); // 打印 d mod (p-1)
	printf("\n");
	BN_print_fp(stdout, RSA_get0_dmq1(rsa)); // 打印 d mod (q-1)
	printf("\n");
	BN_print_fp(stdout, RSA_get0_iqmp(rsa)); // 打印 q^-1 mod p
	printf("\n");
}
void RsaTest(TEEC_Session *sess)
{
	BIGNUM *rsa_key_parts[RSA_KEY_PART_NUM];
	int rc = 0;
	RSA *rsa = RSA_new();
	AddClient(sess, ".");
	GetRSA(sess);
	// BN_bin2bn()
	printf("Conver RSA key to BIGNUM\n");
	for (size_t i = 0; i < RSA_KEY_PART_NUM; i++)
	{
		rsa_key_parts[i] = BN_bin2bn(rsa_key[i], rsa_key_size[i], NULL);
		BN_print_fp(stdout, rsa_key_parts[i]); // 打印 q^-1 mod p
		printf("\n");
	}

	rc = RSA_set0_key(rsa, rsa_key_parts[RSA_KEY_PART_N], rsa_key_parts[RSA_KEY_PART_E], rsa_key_parts[RSA_KEY_PART_D]);
	if (rc == 0)
		errx(1, "RSA_set0_key error %d", rc);
	rc = RSA_set0_factors(rsa, rsa_key_parts[RSA_KEY_PART_P], rsa_key_parts[RSA_KEY_PART_Q]);
	if (rc == 0)
		errx(1, "RSA_set0_factors error %d", rc);
	rc = RSA_set0_crt_params(rsa, rsa_key_parts[RSA_KEY_PART_DMP1], rsa_key_parts[RSA_KEY_PART_DMQ1], rsa_key_parts[RSA_KEY_PART_IQMP]);
	if (rc == 0)
		errx(1, "RSA_set0_crt_params error %d", rc);
	print_rsa_key(rsa);

	char data[] = "The message will be encrypted.";
	char data_encrypted[4096];
	char data_decrypted[4096];
	int data_encrypted_size = RSA_public_encrypt(sizeof(data), data, data_encrypted, rsa, RSA_PKCS1_PADDING);
	printf("encrypted size : %d\n", data_encrypted_size);
	int data_decrypted_size = RSA_private_decrypt(data_encrypted_size, data_encrypted, data_decrypted, rsa, RSA_PKCS1_PADDING);
	printf("decrypted size : %d\n", data_decrypted_size);
	printf("decrypted data : %s\n",data_decrypted);

	char data2[] = "The message will be encrypted too.";
	data_encrypted_size = RSA_private_encrypt(sizeof(data2), data2, data_encrypted, rsa, RSA_PKCS1_PADDING);
	printf("encrypted size : %d\n", data_encrypted_size);
	data_decrypted_size = RSA_public_decrypt(data_encrypted_size, data_encrypted, data_decrypted, rsa, RSA_PKCS1_PADDING);
	printf("decrypted size : %d\n", data_decrypted_size);
	

	printf("decrypted data : %s\n",data_decrypted);
	RSA_free(rsa);
}

void test()
{
	short int x;

	char x0, x1;

	x = 0x1122;

	x0 = ((char *)&x)[0]; // 低地址单元
	x1 = ((char *)&x)[1]; // 高地址单元

	if (0x11 == x0 && 0x22 == x1)
	{
		printf("Big_endian\n");
	}
	else
	{
		printf("Little_endian\n");
	}
}
int main(int argc, char *argv[])
{
	test();
	TEEC_Result res;
	TEEC_Context ctx;
	TEEC_Session sess;
	// TEEC_Operation op;
	TEEC_UUID uuid = TA_API_TEST_UUID;
	uint32_t err_origin;

	/* Initialize a context connecting us to the TEE */
	res = TEEC_InitializeContext(NULL, &ctx);
	if (res != TEEC_SUCCESS)
		errx(1, "TEEC_InitializeContext failed with code 0x%x", res);

	/*
	 * Open a session to the "hello world" TA, the TA will print "hello
	 * world!" in the log when the session is created.
	 */
	res = TEEC_OpenSession(&ctx, &sess, &uuid,
						   0xC0000000, NULL, NULL, &err_origin);
	if (res != TEEC_SUCCESS)
		errx(1, "TEEC_Opensession failed with code 0x%x origin 0x%x",
			 res, err_origin);

	/*
	 * Execute a function in the TA by invoking it, in this case
	 * we're incrementing a number.
	 *
	 * The value of command ID part and how the parameters are
	 * interpreted is part of the interface provided by the TA.
	 */

	/* Clear the TEEC_Operation struct */
	// memset(&op, 0, sizeof(op));

	// op.paramTypes = TEEC_PARAM_TYPES(TEEC_MEMREF_TEMP_INPUT, TEEC_NONE,
	// 								 TEEC_NONE, TEEC_NONE);

	// printf("Invoking TA to increment %d\n", op.params[0].value.a);
	// res = TEEC_InvokeCommand(&sess, TA_API_TEST_CMD_TEST, &op, &err_origin);
	// if (res != TEEC_SUCCESS)
	// 	errx(1, "TEEC_InvokeCommand failed with code 0x%x origin 0x%x", res, err_origin);

	for (size_t i = 0; i < argc; i++)
	{
		printf("ARG[%ld] : %s\n", i, argv[i]);
		if (strcmp(argv[i], "-a") == 0 && i < argc - 1)
		{
			i++;
			AddClient(&sess, argv[i]);
		}
		if (strcmp(argv[i], "-d") == 0 && i < argc - 1)
		{
			i++;
			DeleteClient(&sess, argv[i]);
		}
		if (strcmp(argv[i], "-g") == 0)
		{
			GetRSA(&sess);
		}
		if (strcmp(argv[i], "-t") == 0)
		{
			RsaTest(&sess);
		}
	}

	/*
	 * We're done with the TA, close the session and
	 * destroy the context.
	 *
	 * The TA will print "Goodbye!" in the log when the
	 * session is closed.
	 */

	TEEC_CloseSession(&sess);

	TEEC_FinalizeContext(&ctx);

	return 0;
}
